{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<diskcache.core.Cache at 0x1f1995fc8f0>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np # type: ignore\n",
    "from numba import njit # type: ignore\n",
    "import pandas as pd \n",
    "import pybroker\n",
    "from pybroker import Strategy\n",
    "import talib  \n",
    "import matplotlib.pyplot as plt # type: ignore\n",
    "import akshare as ak #提取上证50股票代码\n",
    "from pybroker.ext.data import AKShare # type: ignore\n",
    "pybroker.enable_data_source_cache('AKShare')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 创建自定义均线函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ma(bar_data,short_window,long_window):\n",
    "    #@njit  # Enable Numba JIT.\n",
    "    def vec_ma(values):\n",
    "        # Initialize the result array.\n",
    "        n = len(values)\n",
    "        \n",
    "        data = np.full((n, 2), np.nan)  \n",
    "        # 将NumPy数组转换为pandas DataFrame，并添加列标签  \n",
    "        df = pd.DataFrame(data, columns=['SMA_short', 'SMA_long'])  \n",
    "\n",
    "        # 计算短期（例如5日）和长期（例如20日）均线  \n",
    "        df['SMA_short'] = talib.SMA(values, timeperiod=short_window)  \n",
    "        df['SMA_long'] = talib.SMA(values, timeperiod=long_window)  \n",
    "            # 初始化金叉和死叉的列  \n",
    "        df['Cross'] = 0  \n",
    "        out = np.array([np.nan for _ in range(n)])\n",
    "\n",
    "        # For all bars starting at lookback:\n",
    "        # 遍历DataFrame，判断金叉和死叉  \n",
    "        for i in range(1, len(df)):  \n",
    "            if df['SMA_short'].iloc[i] > df['SMA_long'].iloc[i] and df['SMA_short'].iloc[i-1] < df['SMA_long'].iloc[i-1]:  \n",
    "                #df.at[df.index[i], 'Cross'] = 2 \n",
    "                out[i] =1\n",
    "            if df['SMA_short'].iloc[i] < df['SMA_long'].iloc[i] and df['SMA_short'].iloc[i-1] > df['SMA_long'].iloc[i-1]:  \n",
    "                #df.at[df.index[i], 'Cross'] = 1  \n",
    "                out[i] = -1\n",
    "        return out\n",
    "    return vec_ma(bar_data.close)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 把自定义函数注册到框架内部"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'指标函数中在 bar_data 之后的任何参数都将作为用户定义的参数传递给 pybroker.indicator。\\n一旦指标函数在 PyBroker 中注册，它将返回一个新的 Indicator 实例，该实例引用我们定义的指标函数。'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 我们将指标函数命名为 cmma_20，并将 回溯 参数指定为 20 条。\n",
    "ma_520 = pybroker.indicator('ma_520', ma, short_window=5, long_window=20)\n",
    "'''指标函数中在 bar_data 之后的任何参数都将作为用户定义的参数传递给 pybroker.indicator。\n",
    "一旦指标函数在 PyBroker 中注册，它将返回一个新的 Indicator 实例，该实例引用我们定义的指标函数。'''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 创建买卖策略函数\n",
    "def buy_ma_cross(ctx):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def buy_ma_cross(ctx):\n",
    "    if ctx.long_pos():\n",
    "        return\n",
    "    # Place a buy order if the most recent value of the 20 day CMMA is < 0:\n",
    "    # 这里调用了 ctx 对象的 indicator 方法，并传入了一个20日CMMA（自定义移动平均差异）指标的名字。\n",
    "    # [-1] 表示获取这个指标数组的最后一个值（即最新的值）。\n",
    "    # 如果这个值小于0，那么说明某种交易条件被触发了（可能是CMMA线向下穿越了零线）。\n",
    "    if ctx.indicator('ma_520')[-1] > 0:\n",
    "        #ctx.buy_shares = ctx.calc_target_shares(1000)\n",
    "        ctx.buy_shares=100 #金叉买\n",
    "    if ctx.indicator('ma_520')[-1] < 0:    \n",
    "        #ctx.hold_bars = 5 # 表示策略计划持有这些股票3个交易日\n",
    "        ctx.sell_shares = 100 #死叉全卖"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 创建上证50列表"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['688981', '688041', '601988', '601601', '600150', '688111', '601728', '601390', '600089', '600050', '688599', '601669', '601225', '600406', '600111', '600905', '603799', '601919', '600900', '601633', '600436', '600438', '600809', '600893', '601899', '603288', '603501', '603986', '601012', '600031', '603259', '601888', '600276', '600690', '600309', '600887', '601288', '601668', '601857', '601088', '601318', '601166', '601628', '600048', '601398', '600519', '600030', '600104', '600036', '600028']\n"
     ]
    }
   ],
   "source": [
    "\n",
    "index_stock_cons_df = ak.index_stock_cons(symbol=\"000016\")\n",
    "\n",
    "# 创建一个空列表，用于存储股票代码\n",
    "symbols = []\n",
    "for stock in index_stock_cons_df['品种代码']:\n",
    "    symbols.append(stock)\n",
    "print(symbols)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 运行回测"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Backtesting: 2023-04-01 00:00:00 to 2024-06-07 00:00:00\n",
      "\n",
      "Loaded cached bar data.\n",
      "\n",
      "Loaded cached indicator data.\n",
      "\n",
      "Test split: 2023-04-03 00:00:00 to 2024-06-07 00:00:00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0% (0 of 286) |                        | Elapsed Time: 0:00:00 ETA:  --:--:--\n",
      "  7% (21 of 286) |#                      | Elapsed Time: 0:00:00 ETA:   0:00:00\n",
      " 10% (31 of 286) |##                     | Elapsed Time: 0:00:00 ETA:   0:00:02\n",
      " 14% (41 of 286) |###                    | Elapsed Time: 0:00:00 ETA:   0:00:04\n",
      " 17% (51 of 286) |####                   | Elapsed Time: 0:00:01 ETA:   0:00:06\n",
      " 21% (61 of 286) |####                   | Elapsed Time: 0:00:01 ETA:   0:00:06\n",
      " 24% (71 of 286) |#####                  | Elapsed Time: 0:00:02 ETA:   0:00:06\n",
      " 28% (81 of 286) |######                 | Elapsed Time: 0:00:02 ETA:   0:00:06\n",
      " 31% (91 of 286) |#######                | Elapsed Time: 0:00:02 ETA:   0:00:06\n",
      " 35% (101 of 286) |#######               | Elapsed Time: 0:00:03 ETA:   0:00:06\n",
      " 38% (111 of 286) |########              | Elapsed Time: 0:00:03 ETA:   0:00:06\n",
      " 42% (121 of 286) |#########             | Elapsed Time: 0:00:04 ETA:   0:00:05\n",
      " 45% (131 of 286) |##########            | Elapsed Time: 0:00:04 ETA:   0:00:05\n",
      " 49% (141 of 286) |##########            | Elapsed Time: 0:00:05 ETA:   0:00:05\n",
      " 52% (151 of 286) |###########           | Elapsed Time: 0:00:05 ETA:   0:00:05\n",
      " 56% (161 of 286) |############          | Elapsed Time: 0:00:06 ETA:   0:00:04\n",
      " 59% (171 of 286) |#############         | Elapsed Time: 0:00:06 ETA:   0:00:04\n",
      " 63% (181 of 286) |#############         | Elapsed Time: 0:00:07 ETA:   0:00:04\n",
      " 66% (191 of 286) |##############        | Elapsed Time: 0:00:07 ETA:   0:00:03\n",
      " 70% (201 of 286) |###############       | Elapsed Time: 0:00:08 ETA:   0:00:03\n",
      " 73% (211 of 286) |################      | Elapsed Time: 0:00:08 ETA:   0:00:02\n",
      " 77% (221 of 286) |#################     | Elapsed Time: 0:00:08 ETA:   0:00:02\n",
      " 80% (231 of 286) |#################     | Elapsed Time: 0:00:08 ETA:   0:00:02\n",
      " 84% (241 of 286) |##################    | Elapsed Time: 0:00:08 ETA:   0:00:01\n",
      " 87% (251 of 286) |###################   | Elapsed Time: 0:00:09 ETA:   0:00:01\n",
      " 91% (261 of 286) |####################  | Elapsed Time: 0:00:09 ETA:   0:00:00\n",
      " 94% (271 of 286) |####################  | Elapsed Time: 0:00:09 ETA:   0:00:00\n",
      " 98% (281 of 286) |##################### | Elapsed Time: 0:00:10 ETA:   0:00:00\n",
      "100% (286 of 286) |######################| Elapsed Time: 0:00:10 ETA:  00:00:00\n",
      "100% (286 of 286) |######################| Elapsed Time: 0:00:10 Time:  0:00:10\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Finished backtest: 0:00:21\n",
      "                      name        value\n",
      "0              trade_count     381.0000\n",
      "1     initial_market_value  100000.0000\n",
      "2         end_market_value   99652.3800\n",
      "3                total_pnl     321.0000\n",
      "4           unrealized_pnl    -668.6200\n",
      "5         total_return_pct       0.3210\n",
      "6             total_profit   47914.0000\n",
      "7               total_loss  -47593.0000\n",
      "8               total_fees       0.0000\n",
      "9             max_drawdown  -18919.4900\n",
      "10        max_drawdown_pct     -16.1820\n",
      "11                win_rate      37.9947\n",
      "12               loss_rate      62.0053\n",
      "13          winning_trades     144.0000\n",
      "14           losing_trades     235.0000\n",
      "15                 avg_pnl       0.8425\n",
      "16          avg_return_pct       0.0328\n",
      "17          avg_trade_bars      16.1549\n",
      "18              avg_profit     332.7361\n",
      "19          avg_profit_pct       6.7906\n",
      "20  avg_winning_trade_bars      26.5833\n",
      "21                avg_loss    -202.5234\n",
      "22            avg_loss_pct      -4.1078\n",
      "23   avg_losing_trade_bars       9.7362\n",
      "24             largest_win    6569.0000\n",
      "25         largest_win_pct      26.9100\n",
      "26        largest_win_bars      35.0000\n",
      "27            largest_loss   -3596.0000\n",
      "28        largest_loss_pct     -10.5800\n",
      "29       largest_loss_bars      10.0000\n",
      "30                max_wins      17.0000\n",
      "31              max_losses      21.0000\n",
      "32                  sharpe      -0.0009\n",
      "33                 sortino      -0.0013\n",
      "34           profit_factor       0.9976\n",
      "35             ulcer_index       3.7088\n",
      "36                     upi       0.0021\n",
      "37               equity_r2       0.3112\n",
      "38               std_error    4474.0171\n"
     ]
    }
   ],
   "source": [
    "aKShare = AKShare()\n",
    "strategy = Strategy(aKShare, '4/1/2023', '6/7/2024')\n",
    "strategy.add_execution(buy_ma_cross,  symbols, indicators=ma_520)\n",
    "pybroker.enable_indicator_cache('my_indicators')\n",
    "\n",
    "result = strategy.backtest(warmup=20) #warmup 参数指定在运行回测执行之前需要经过 20 个 Bar：\n",
    "# 在Python的pandas库中，DataFrame 对象有一个方法 round()，\n",
    "# 该方法用于将DataFrame中的所有数值型列（通常是整数或浮点数）四舍五入到指定的小数位数。\n",
    "result.metrics_df.round(4)\n",
    "print(result.metrics_df.round(4))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
